#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "platform.h"
#include "xil_printf.h"
#include "xparameters.h"
#include "xuartps.h"
#include "xil_cache.h"
// 删除 ff.h 引用，因为我们不需要文件系统功能
#include "xil_io.h"
#include "sleep.h"

// MNIST模型配置
#define IMAGE_SIZE 28
#define IMAGE_BYTES (IMAGE_SIZE * IMAGE_SIZE)
#define MAX_BUFFER_SIZE 1024

// BRAM地址配置 (根据您的Block Design调整)
#define BRAM_BASE_ADDR XPAR_AXI_BRAM_CTRL_0_S_AXI_BASEADDR
#define RESULT_OFFSET  0x1000

// UART配置
#define UART_DEVICE_ID XPAR_XUARTPS_0_DEVICE_ID

// CNN网络权重和参数
// 注：实际项目中这些应该从内存加载或存储在代码中
static float conv1_weights[16][1][5][5];  // 第一层卷积权重
static float conv1_bias[16];              // 第一层卷积偏置
static float fc1_weights[10][256];        // 全连接层权重
static float fc1_bias[10];                // 全连接层偏置

// UART实例
XUartPs UartPs;

// 初始化UART
int init_uart() {
    XUartPs_Config *Config;
    int Status;
    
    Config = XUartPs_LookupConfig(UART_DEVICE_ID);
    if (NULL == Config) {
        return XST_FAILURE;
    }
    
    Status = XUartPs_CfgInitialize(&UartPs, Config, Config->BaseAddress);
    if (Status != XST_SUCCESS) {
        return XST_FAILURE;
    }
    
    XUartPs_SetBaudRate(&UartPs, 115200);
    
    return XST_SUCCESS;
}

// 接收图像数据
int receive_image(uint8_t *image_buffer) {
    char buffer[MAX_BUFFER_SIZE];
    int bytes_received = 0;
    int in_data_mode = 0;
    
    while (1) {
        int received = XUartPs_Recv(&UartPs, buffer, MAX_BUFFER_SIZE);
        
        if (received > 0) {
            if (!in_data_mode) {
                // 检查开始标记
                if (strstr(buffer, "START") != NULL) {
                    in_data_mode = 1;
                    xil_printf("开始接收图像数据...\n\r");
                    continue;
                }
            } else {
                // 检查结束标记
                if (strstr(buffer, "END") != NULL) {
                    xil_printf("图像数据接收完成\n\r");
                    break;
                }
                
                // 复制数据到图像缓冲区
                if (bytes_received + received <= IMAGE_BYTES) {
                    memcpy(image_buffer + bytes_received, buffer, received);
                    bytes_received += received;
                }
            }
        }
        usleep(1000);
    }
    
    return bytes_received;
}

// 将图像数据写入BRAM
void write_image_to_bram(uint8_t *image_buffer) {
    for (int i = 0; i < IMAGE_BYTES; i += 4) {
        // 每次写入32位(4字节)
        uint32_t word = 0;
        for (int j = 0; j < 4 && (i + j) < IMAGE_BYTES; j++) {
            word |= (image_buffer[i + j] << (j * 8));
        }
        Xil_Out32(BRAM_BASE_ADDR + i, word);
    }
    
    // 刷新缓存，确保数据写入BRAM
    Xil_DCacheFlush();
}

// 简化版CNN推理 (软件实现，实际项目中应使用硬件加速)
int run_mnist_inference(uint8_t *image_data) {
    // 简化的CNN实现
    // 实际项目中，这部分应该由硬件加速器完成
    
    // 预处理
    float input[28][28];
    for (int i = 0; i < 28; i++) {
        for (int j = 0; j < 28; j++) {
            input[i][j] = image_data[i * 28 + j] / 255.0f;
        }
    }
    
    // 这里添加CNN前向传播代码
    // ... 卷积层，池化层，全连接层等
    
    // 返回结果 (0-9)
    // 模拟一个结果
    int result = image_data[0] % 10;  // 仅用于演示
    
    return result;
}

// 发送识别结果到PC
void send_result(int result) {
    char buffer[32];
    sprintf(buffer, "RESULT:%d\r\n", result);
    XUartPs_Send(&UartPs, (u8*)buffer, strlen(buffer));
}

int main() {
    init_platform();
    
    // 初始化UART
    if (init_uart() != XST_SUCCESS) {
        xil_printf("UART初始化失败\r\n");
        return 1;
    }
    
    xil_printf("MNIST识别系统已启动\r\n");
    xil_printf("等待PC发送图像数据...\r\n");
    
    // 分配图像数据缓冲区
    uint8_t *image_buffer = (uint8_t*)malloc(IMAGE_BYTES);
    if (!image_buffer) {
        xil_printf("内存分配失败\r\n");
        return 1;
    }
    
    while (1) {
        // 接收图像数据
        int bytes_received = receive_image(image_buffer);
        
        if (bytes_received > 0) {
            xil_printf("成功接收 %d 字节数据\r\n", bytes_received);
            
            // 将图像数据写入BRAM
            write_image_to_bram(image_buffer);
            
            // 执行推理
            xil_printf("开始执行推理...\r\n");
            int result = run_mnist_inference(image_buffer);
            
            // 发送结果到PC
            xil_printf("识别结果: %d\r\n", result);
            send_result(result);
        }
    }
    
    free(image_buffer);
    cleanup_platform();
    return 0;
}
